#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
batch_asteroid_calmag_with_util.py

批量检测小行星，提取 CALMAG_22，绘制光变曲线，生成小行星切片马赛克与中心切片 GIF。
"""

import os
import sys
import argparse
from datetime import datetime
from loguru import logger

from astroquery.imcce import Skybot
from astropy.io import fits
from astropy.coordinates import SkyCoord
from astropy.time import Time
import astropy.units as u
from astropy.table import Table
from astropy.wcs import WCS
from astropy.visualization import simple_norm

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.patches import Rectangle

from tools.utils import TargetXCatalogue    # 交叉匹配函数
from tools.make_gif import make_gif         # GIF 生成函数
import ntpath


def find_asteroids(fits_path, radius_arcmin=22.0):
    """
    在一张 FITS 图像中使用 Skybot 搜索已知小行星。

    返回值：
        obs_time : datetime     FITS 中 DATE 字段对应的观测时间
        result   : astropy.Table 小行星列表
    """
    hdr     = fits.getheader(fits_path, 0)
    # 从 FITS 头读取中心 RA/DEC 和 DATE
    ra_hdr  = hdr.get('TEL_RA')   or hdr.get('CRVAL1')
    dec_hdr = hdr.get('TEL_DEC')  or hdr.get('CRVAL2')
    date    = hdr.get('DATE')     or hdr.get('DATE-OBS')
    if ra_hdr is None or dec_hdr is None or date is None:
        raise KeyError("无法从 FITS 头中获取 TEL_RA/TEL_DEC/DATE")

    coord = SkyCoord(float(ra_hdr)*u.deg, float(dec_hdr)*u.deg)
    epoch = Time(date)
    # 返回观测时间 + Skybot 搜索结果
    return epoch.to_datetime(), Skybot.cone_search(
        coord,
        radius_arcmin * u.arcmin,
        epoch,
        'O44'
    )


def load_calmag_catalog(cat_path):
    """
    载入 SExtractor 输出的星表，并检查必要列是否存在。
    """
    with fits.open(cat_path, mode='update') as hdul:
        hdr = hdul[0].header
        # 删除损坏关键词
        for key in ["TTYPE99", "TFORM99", "TUNIT99"]:
            if key in hdr:
                del hdr[key]
        hdul.flush()
        
    tbl = Table.read(cat_path, hdu=2)
    for col in ('ALPHA_J2000','DELTA_J2000','CALMAG_22','MAGERR_AUTO_S'):
        if col not in tbl.colnames:
            raise KeyError(f"{cat_path} 缺少列 {col}")
    return tbl


def crossmatch_with_util(ast_tbl, phot_tbl, max_sep_arcsec=2.0):
    """
    使用 TargetXCatalogue 在小行星列表和星表之间做交叉匹配。
    返回每颗小行星的 (CALMAG, MAGERR)。
    """
    results = []
    for row in ast_tbl:
        ra, dec = row['RA'].value, row['DEC'].value
        match = TargetXCatalogue(
            ra, dec, phot_tbl,
            radius=max_sep_arcsec,
            RaKeyword='ALPHA_J2000',
            DecKeyword='DELTA_J2000'
        )
        if match is not None:
            calmag = match['CALMAG_22']
            magerr = match['MAGERR_AUTO_S']
            flags = match['FLAGS']
        else:
            calmag = magerr = None
        results.append((calmag, magerr,flags))
    return results


def process_directory(root_dir, radius=22.0, max_sep=2.0):
    """
    遍历根目录，处理每个 *sciimg.fits：
    - 搜索小行星
    - 载入对应的星表
    - 交叉匹配提取 CALMAG & MAGERR
    汇总所有波段的数据返回。
    """
    all_data = []  # (obs_time, band, calmag, magerr)
    all_data_withradec=[]
    for dirpath, _, files in os.walk(root_dir):
        for fn in files:
            if not fn.lower().endswith('sciimg.fits'):
                continue
            # 从文件名中解析滤波器波段（例如 mr_sc_<object>_i_yyyymmdd...）
            parts = fn.split('_')
            band  = parts[3] if len(parts) > 3 else 'unknown'
            path  = os.path.join(dirpath, fn)

            # 搜索小行星
            try:
                obs_time, ast_tbl = find_asteroids(path, radius)
                # logger.info(f"ast_tbl:{ast_tbl}")
                ra_, dec_ = ast_tbl[0]['RA'].value, ast_tbl[0]['DEC'].value
            except Exception as e:
                logger.error(f"{fn}: 小行星检测失败：{e}")
                continue

            # 对应的星表文件
            base     = fn[:-5]
            cat_path = os.path.join(dirpath, f"{base}_sexcat_CalMag.fits")
            if not os.path.exists(cat_path):
                continue

            # 载入星表 & 交叉匹配
            phot_tbl = load_calmag_catalog(cat_path)
            matches  = crossmatch_with_util(ast_tbl, phot_tbl, max_sep)
            for calmag, magerr,flags in matches:
                if calmag is not None:
                    
                    all_data.append((obs_time, band, calmag, magerr))
                    all_data_withradec.append((fn,obs_time, band, calmag, magerr,ra_, dec_,flags))
    return all_data,all_data_withradec


def plot_time_mag(data, separate=False):
    """
    绘制光变曲线（CALMAG vs 时间）：
      - separate=False: 所有波段合并一图
      - separate=True: 每个波段单独子图
    时间轴使用小时主刻度、15 分钟次刻度。
    """
    from matplotlib.dates import DateFormatter, HourLocator, MinuteLocator

    bands = sorted({d[1] for d in data})

    def format_axes(ax):
        ax.invert_yaxis()
        ax.set_xlabel('Observation Time')
        ax.set_ylabel('CALMAG_22')
        ax.xaxis.set_major_locator(HourLocator(interval=1))
        ax.xaxis.set_major_formatter(DateFormatter('%m-%d %H:%M'))
        ax.xaxis.set_minor_locator(MinuteLocator(byminute=[0,15,30,45]))
        ax.grid(True, which='both', axis='x', linestyle='--', alpha=0.5)

    # 分波段单图
    if separate:
        fig, axes = plt.subplots(len(bands), 1, sharex=True,
                                 figsize=(10, 3*len(bands)))
        if len(bands) == 1:
            axes = [axes]
        for ax, band in zip(axes, bands):
            times = [d[0] for d in data if d[1] == band]
            mags  = [d[2] for d in data if d[1] == band]
            errs  = [d[3] for d in data if d[1] == band]
            ax.errorbar(times, mags, yerr=errs, fmt='o', capsize=3, label=band)
            ax.set_title(f'Band {band}')
            ax.legend()
            format_axes(ax)
        axes[-1].set_xlabel('Observation Time')
        plt.tight_layout()
        plt.savefig("result/time_mag_separate.png", dpi=300)
        plt.close(fig)

    # 合并图
    else:
        fig, ax = plt.subplots(figsize=(12,6))
        for band in bands:
            times = [d[0] for d in data if d[1] == band]
            mags  = [d[2] for d in data if d[1] == band]
            errs  = [d[3] for d in data if d[1] == band]
            ax.errorbar(times, mags, yerr=errs, fmt='o', capsize=3, label=band)
        format_axes(ax)
        ax.set_title('Asteroid CALMAG vs Time by Band')
        ax.legend(title='Filter')
        plt.tight_layout()
        plt.savefig("result/time_mag.png", dpi=300)
        plt.close(fig)


def make_asteroid_mosaic(root_dir, output_path, size=200, max_per_row=5,
                         radius=22.0, max_sep=2.0,
                         band_order=('u', 'g', 'r', 'i', 'z', 'y')):
    """
    遍历 root_dir 下所有 sciimg.fits，
    对每个文件在小行星像素位置截取 size x size 区域，
    并按网格组合成一张大图保存到 output_path。
    每个子图顶部显示：
      <filename>
      <band>-band: mag=XX.XX±YY.YY
    切片先按 band_order 排序，同一波段内再按文件名末尾数字排序。
    """
    import os
    import numpy as np
    from astropy.visualization import simple_norm
    from astropy.wcs import WCS
    from astropy.coordinates import SkyCoord
    import astropy.units as u
    from astropy.io import fits
    import matplotlib.pyplot as plt

    entries = []  # 存放 (band, idx_number, fn, cutout, mag, err)

    for dirpath, _, files in os.walk(root_dir):
        for fn in files:
            if not fn.lower().endswith('sciimg.fits'):
                continue

            parts = fn.split('_')
            band = parts[3] if len(parts) > 3 else 'unknown'
            # 取倒数第二段作为索引数字
            try:
                idx_num = int(parts[-3])
            except ValueError:
                idx_num = float('inf')

            fits_path = os.path.join(dirpath, fn)

            # 1) 找小行星
            _, ast_tbl = find_asteroids(fits_path, radius)
            if not ast_tbl:
                continue

            # 2) 取第一个小行星的 mag/err
            base     = fn[:-5]
            cat_path = os.path.join(dirpath, f"{base}_sexcat_CalMag.fits")
            if not os.path.exists(cat_path):
                continue
            phot_tbl = load_calmag_catalog(cat_path)
            calmag, magerr,flags = crossmatch_with_util(ast_tbl, phot_tbl, max_sep)[0]

            # 3) 裁剪
            ra, dec = ast_tbl[0]['RA'].value, ast_tbl[0]['DEC'].value
            hdul = fits.open(fits_path)
            data = hdul[0].data
            wcs  = WCS(hdul[0].header)
            xpix, ypix = wcs.world_to_pixel(SkyCoord(ra*u.deg, dec*u.deg))
            x, y = int(xpix), int(ypix)
            half = size // 2
            cut  = data[y-half:y+half, x-half:x+half]

            entries.append((band, idx_num, fn, cut, calmag, magerr))

    if not entries:
        return

    # 先按 band_order 排序，再按 idx_num
    def sort_key(entry):
        b, idx, *_ = entry
        try:
            bidx = band_order.index(b)
        except ValueError:
            bidx = len(band_order)
        return (bidx, idx)

    entries.sort(key=sort_key)

    n    = len(entries)
    ncol = min(max_per_row, n)
    nrow = int(np.ceil(n / ncol))

    fig, axes = plt.subplots(nrow, ncol, figsize=(2*ncol, 2*nrow))
    axes = axes.flatten()

    # 隐藏多余子图
    for ax in axes[n:]:
        ax.axis('off')

    # 绘制
    for i, (_, idx, fn, cut, mag, err) in enumerate(entries):
        ax = axes[i]
        norm = simple_norm(cut, 'sqrt', percent=99)
        ax.imshow(cut, norm=norm, cmap='gray', origin='lower')
        ax.set_title(f"{fn}\n#{idx}: {mag:.2f}±{err:.2f}", fontsize=3)
        ax.axis('off')

    plt.tight_layout()
    plt.savefig(output_path, dpi=300)
    plt.close(fig)




def make_region_cutout(fits_path, x_pix, y_pix, output_dir, size=500,filiename=None):
    """
    从单帧 FITS 中按给定像素位置裁剪 size×size 区域，
    并在中心位置画红框标记小行星，输出 PNG。
    """
    data = fits.getdata(fits_path)
    half = size // 2
    cut  = data[y_pix-half:y_pix+half, x_pix-half:x_pix+half]
    norm = simple_norm(cut, 'sqrt', percent=99)

    fig, ax = plt.subplots(figsize=(5,5), dpi=100)
    ax.imshow(cut, norm=norm, cmap='gray', origin='lower')

    # 在中心位置绘制 20×20 的红框
    rect = Rectangle((half-10, half-10), 20, 20,
                     edgecolor='red', fill=False, lw=1.5)
    ax.add_patch(rect)
    ax.axis('off')

    os.makedirs(output_dir, exist_ok=True)
    outfn = os.path.splitext(os.path.basename(fits_path))[0] + '.png'
    fig.savefig(os.path.join(output_dir, outfn),
                bbox_inches='tight', pad_inches=0)
    plt.close(fig)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="批量检测小行星、提取 CALMAG、绘制光变曲线、生成切片与马赛克/GIF"
    )
    parser.add_argument("-d", "--directory",
                        default="./images/12377/",
                        help="根目录，遍历 sciimg.fits")
    parser.add_argument("-r", "--radius", type=float, default=22.0,
                        help="Skybot 搜索半径（arcmin）")
    parser.add_argument("-s", "--maxsep", type=float, default=2.0,
                        help="交叉匹配最大分离（arcsec）")
    parser.add_argument("--separate", action="store_true",
                        help="为各波段单独子图")
    parser.add_argument("--mosaic", default="True",
                        help="生成小行星切片马赛克")
    parser.add_argument("--cutout", default="True",
                        help="生成每帧中心切片并标记")
    parser.add_argument("--size", type=int, default=200,
                        help="切片/马赛克单元大小（px）")
    parser.add_argument("--maxcol", type=int, default=5,
                        help="马赛克最大列数")
    parser.add_argument("--fps", type=float, default=2.0,
                        help="中心切片 GIF 帧率 (frame/sec)")
    args = parser.parse_args()

    # 依赖检查
    try:
        import astroquery  # noqa
        import astropy     # noqa
        import matplotlib  # noqa
    except ImportError as ie:
        sys.exit(f"缺少依赖：{ie.name}，请安装 astroquery astropy matplotlib")

    os.makedirs("result", exist_ok=True)

    # 1. 绘制光变曲线
    data,datawithradec = process_directory(args.directory, args.radius, args.maxsep)
    if data:
        plot_time_mag(data, args.separate)
        logger.info("光变曲线图输出到 result/")
        import pandas as pd
        df = pd.DataFrame(datawithradec, columns=["fn","obs_time", "band", "calmag", "magerr","ra","dec","flags"])
        # 保存为 CSV
        df.to_csv("result/asteroid_photometry_results.csv", index=False)

    # 2. 生成马赛克
    if args.mosaic:
        outm = "result/mosaic.png"
        make_asteroid_mosaic(args.directory, outm,size=args.size, max_per_row=args.maxcol,band_order=('u','g','r','i','z'))
        logger.info(f"马赛克图输出到 {outm}")

    # 3. 生成中心切片并合成 GIF
    if args.cutout:
        for dirpath, _, files in os.walk(args.directory):
            for fn in files:
                if not fn.lower().endswith("sciimg.fits"):
                    continue
                path = os.path.join(dirpath, fn)
                _, ast_tbl = find_asteroids(path, args.radius)
                if not ast_tbl:
                    continue

                # 对应 FITS 中小行星第一颗的位置
                ra, dec = ast_tbl[0]['RA'].value, ast_tbl[0]['DEC'].value
                wcs = WCS(fits.getheader(path))
                xpix, ypix = wcs.world_to_pixel(SkyCoord(ra*u.deg, dec*u.deg))

                # 按滤波器分目录保存
                filt  = fits.getheader(path)["FILTER"]
                cutdir = f"result/cutouts/{filt}"
                make_region_cutout(path, int(xpix), int(ypix),cutdir, size=500)
            
        # 生成 GIF：请带上后缀 .gif
        for filt in ["u", "v", "g", "r","i","z"]:
            cutdir = f"result/cutouts/{filt}"
            if os.path.join(cutdir):
                gif_path = f"result/asteroid_{filt}.gif"
                make_gif(cutdir, gif_path, pattern="*.png", fps=args.fps)
                logger.info(f"中心切片 GIF 输出到 {gif_path}")
